Feat: sebulba rec ippo #1142
Open
Feat: sebulba rec ippo #1142
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
SimonDuToit
requested review from
RuanJohn,
sash-a,
OmaymaMahjoub,
WiemKhlifi and
Louay-Ben-nessir
as code owners
OmaymaMahjoub
requested changes
Dec 11, 2024
| @@ -0,0 +1,910 @@ | |||
| # Copyright 2022 InstaDeep Ltd. All rights reserved. | |||
There was a problem hiding this comment.
If you can update the typings in Pipeline mava/utils/sebulba.py to be Union[PPOTransition, RNNPPOTransition]
There was a problem hiding this comment.
This causes errors in the pre-commit. For now I changed both sebulba systems to use the MavaTransition type-var but this is probably a temporary solution.
There was a problem hiding this comment.
Can you please make an issue for this. I think the best solution is to make a protocol with the all the common things in a transition (actions, obs, done, reward). The challenge is that named tuples don't seem to work with protocols so we'd likely need to switch to a flax/chex dataclass
| log_prob = actor_policy.log_prob(action) | ||
| # It may be faster to calculate the values in the learner as | ||
| # then we won't need to pass critic params to actors. | ||
| # value = critic_apply_fn(params.critic_params, observation).squeeze() |
There was a problem hiding this comment.
if you can remove this comment
sash-a
requested changes
Jan 9, 2025
|
|
||
| def _update_minibatch(train_state: Tuple, batch_info: Tuple) -> Tuple: | ||
| """Update the network for a single minibatch.""" | ||
| # UNPACK TRAIN STATE AND BATCH INFO |
There was a problem hiding this comment.
In line with PPO updates I did last year
Suggested change
| # UNPACK TRAIN STATE AND BATCH INFO |
| key: chex.PRNGKey, | ||
| ) -> Tuple: | ||
| """Calculate the actor loss.""" | ||
| # RERUN NETWORK |
There was a problem hiding this comment.
Suggested change
| # RERUN NETWORK | |
| # Rerun network |
| targets: chex.Array, | ||
| ) -> Tuple: | ||
| """Calculate the critic loss.""" | ||
| # RERUN NETWORK |
There was a problem hiding this comment.
Suggested change
| # RERUN NETWORK | |
| # Rerun network |
| critic_params, traj_batch.hstates.critic_hidden_state[0], obs_and_done | ||
| ) | ||
|
|
||
| # CALCULATE VALUE LOSS |
There was a problem hiding this comment.
Suggested change
| # CALCULATE VALUE LOSS | |
| # Calculate value loss |
Comment on lines
+305
to
+320
| loss_actor1 = ratio * gae | ||
| loss_actor2 = ( | ||
| jnp.clip( | ||
| ratio, | ||
| 1.0 - config.system.clip_eps, | ||
| 1.0 + config.system.clip_eps, | ||
| ) | ||
| * gae | ||
| ) | ||
| loss_actor = -jnp.minimum(loss_actor1, loss_actor2) | ||
| loss_actor = loss_actor.mean() | ||
| # The seed will be used in the TanhTransformedDistribution: | ||
| entropy = actor_policy.entropy(seed=key).mean() | ||
|
|
||
| total_loss = loss_actor - config.system.ent_coef * entropy | ||
| return total_loss, (loss_actor, entropy) |
There was a problem hiding this comment.
Suggested change
| loss_actor1 = ratio * gae | |
| loss_actor2 = ( | |
| jnp.clip( | |
| ratio, | |
| 1.0 - config.system.clip_eps, | |
| 1.0 + config.system.clip_eps, | |
| ) | |
| * gae | |
| ) | |
| loss_actor = -jnp.minimum(loss_actor1, loss_actor2) | |
| loss_actor = loss_actor.mean() | |
| # The seed will be used in the TanhTransformedDistribution: | |
| entropy = actor_policy.entropy(seed=key).mean() | |
| total_loss = loss_actor - config.system.ent_coef * entropy | |
| return total_loss, (loss_actor, entropy) | |
| actor_loss1 = ratio * gae | |
| actor_loss2 = ( | |
| jnp.clip( | |
| ratio, | |
| 1.0 - config.system.clip_eps, | |
| 1.0 + config.system.clip_eps, | |
| ) | |
| * gae | |
| ) | |
| actor_loss = -jnp.minimum(actor_loss1, actor_loss2) | |
| actor_loss = actor_loss.mean() | |
| # The seed will be used in the TanhTransformedDistribution: | |
| entropy = actor_policy.entropy(seed=key).mean() | |
| total_loss = actor_loss - config.system.ent_coef * entropy | |
| return total_loss, (actor_loss, entropy) |
|
|
||
| # Calculate critic loss | ||
| critic_grad_fn = jax.value_and_grad(_critic_loss_fn, has_aux=True) | ||
| critic_loss_info, critic_grads = critic_grad_fn( |
There was a problem hiding this comment.
Suggested change
| critic_loss_info, critic_grads = critic_grad_fn( | |
| value_loss_info, critic_grads = critic_grad_fn( |
Comment on lines
+368
to
+370
| critic_grads, critic_loss_info = jax.lax.pmean( | ||
| (critic_grads, critic_loss_info), axis_name="learner_devices" | ||
| ) |
There was a problem hiding this comment.
Suggested change
| critic_grads, critic_loss_info = jax.lax.pmean( | |
| (critic_grads, critic_loss_info), axis_name="learner_devices" | |
| ) | |
| critic_grads, value_loss_info = jax.lax.pmean( | |
| (critic_grads, value_loss_info), axis_name="learner_devices" | |
| ) |
Comment on lines
+388
to
+396
| actor_total_loss, (actor_loss, entropy) = actor_loss_info | ||
| critic_total_loss, (value_loss) = critic_loss_info | ||
| total_loss = critic_total_loss + actor_total_loss | ||
| loss_info = { | ||
| "total_loss": total_loss, | ||
| "value_loss": value_loss, | ||
| "actor_loss": actor_loss, | ||
| "entropy": entropy, | ||
| } |
There was a problem hiding this comment.
Suggested change
| actor_total_loss, (actor_loss, entropy) = actor_loss_info | |
| critic_total_loss, (value_loss) = critic_loss_info | |
| total_loss = critic_total_loss + actor_total_loss | |
| loss_info = { | |
| "total_loss": total_loss, | |
| "value_loss": value_loss, | |
| "actor_loss": actor_loss, | |
| "entropy": entropy, | |
| } | |
| actor_loss, (_, entropy) = actor_loss_info | |
| value_loss, (unscaled_value_loss) = value_loss_info | |
| total_loss = actor_loss + value_loss | |
| loss_info = { | |
| "total_loss": total_loss, | |
| "value_loss": unscaled_value_loss, | |
| "actor_loss": actor_loss, | |
| "entropy": entropy, | |
| } |
| batch = tree.map( | ||
| lambda x: x.reshape( | ||
| config.system.recurrent_chunk_size, | ||
| num_learner_envs * num_recurrent_chunks, |
There was a problem hiding this comment.
Can you extract this into a variable called batch_size. Will make it more clear and it's also used below
Comment on lines
+406
to
+407
| # batch_size = config.system.rollout_length * num_learner_envs | ||
| # permutation = jax.random.permutation(shuffle_key, batch_size) |
There was a problem hiding this comment.
Suggested change
| # batch_size = config.system.rollout_length * num_learner_envs | |
| # permutation = jax.random.permutation(shuffle_key, batch_size) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Sebulba implementation of recurrent IPPO.